#!/usr/bin/env python3
# -*- coding:utf-8 -*-
####################################################################################################
# File: utils.py                                                                                   #
# Created Date: Sunday, August 15th 2021, 12:21:01 am                                              #
# Author: Zongliang Hou                                                                            #
# E-mail: EnderZ@sjtu.edu.cn                                                                       #
# Copyright (c) 2021 Lei Lab                                                                       #
####################################################################################################


def read_seqfile(file: str) -> str:
    seq = []
    with open(file, 'r') as f:
        for r in f.readlines():
            seq.append(r.strip())
    seq = ''.join(seq)
    return seq


def write_seqfile(seq: str, file: str) -> None:
    with open(file, 'w') as f:
        f.write(seq)


def read_fasta(file) -> dict:
    seqs = {}
    seq = ''
    seq_name = None
    with open(file, 'r') as f:
        for r in f.readlines():
            if r.startswith('>'):
                if seq_name:
                    seqs[seq_name] = seq
                    seq = ''
                seq_name = r[1:].strip()
            else:
                seq = ''.join([seq, r.strip()])
        if seq_name:
            seqs[seq_name] = seq
    return seqs


def complement_dna(sequence: str, reverse: bool = False) -> str:
    """Complement DNA sequence"""
    # Replace A->T and T->A
    sequence = sequence.replace('A', 'N').replace('T', 'A').replace('N', 'T')
    # Replace C->G and G->C
    sequence = sequence.replace('C', 'N').replace('G', 'C').replace('N', 'G')
    # Reversed
    if reverse:
        sequence = sequence[::-1]
    return sequence
